import numpy as np
import cv2
import tritonclient.grpc as grpcclient
import time
import os

def plot_box_label(ori_image, box, label=None, color=(128, 128, 128), txt_color=(255, 255, 255), pil = False, text_lw = 2):
    if pil:
        image = np.asarray(ori_image).copy()
    else:
        image = ori_image
        
    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    cv2.rectangle(image, p1, p2, color, thickness=text_lw, lineType=cv2.LINE_AA)
    if label:
        tf = max(text_lw - 1, 1)  # font thickness
        w, h = cv2.getTextSize(label, 0, fontScale=text_lw / 3, thickness=tf)[0]  # text width, height
        outside = p1[1] - h >= 3
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        
        cv2.rectangle(image, p1, p2, color, thickness=-1, lineType=cv2.LINE_AA)  # filled
        cv2.putText(image,
                    label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                    0,
                    text_lw / 3,
                    txt_color,
                    thickness=tf,
                    lineType=cv2.LINE_AA)
    return np.asarray(image)


if __name__ == '__main__':
    triton_client = grpcclient.InferenceServerClient(url='192.168.96.136:8301')
    score_threshold = 0.3
    input_path = "/workspace/workspace/wumh/wuminghui/12_Smoke_fire/test"
    for img_name in os.listdir(input_path):
        img_path = os.path.join(input_path, img_name)
        image = cv2.imread(img_path)
        img = image.transpose((1, 0, 2))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        if img is None:
            raise FileNotFoundError(f"Image at path {img_path} not found")
        # 设置输入
        inputs = [  
            grpcclient.InferInput('image', [*img.shape], "UINT8"),
            grpcclient.InferInput('score', [1], "FP16")
        ]
        inputs[0].set_data_from_numpy(img)
        inputs[1].set_data_from_numpy(np.array([score_threshold], dtype=np.float16))

        # 设置输出
        outputs = [
            grpcclient.InferRequestedOutput('classes'),
            grpcclient.InferRequestedOutput('scores'),
            grpcclient.InferRequestedOutput('bboxes'),
            grpcclient.InferRequestedOutput("labels")
        ]

        t1 = time.time()
        infer_result = triton_client.infer('base', inputs=inputs, outputs=outputs)
        t2 = time.time()

        # 获取推理结果
        bboxes = infer_result.as_numpy('bboxes')
        scores = infer_result.as_numpy('scores')
        classes = infer_result.as_numpy('classes')
        labels = infer_result.as_numpy('labels')

        for i in range(len(bboxes)):
            print(
                f"label: ['{labels[i].decode('utf-8')}']    class: [{classes[i]}]    score: [{round(scores[i], 4)}]"
                f"    bbox: {bboxes[i]}")
        print('inference time is: {}ms'.format(1000 * (t2 - t1)))

        # 绘图并保存
        img_bgr = image
        for i, box in enumerate(bboxes):
            img_bgr = plot_box_label(
                ori_image=image,
                box=box,
                label=f"id: {classes[i]} {scores[i]:.2f}"
            )
        cv2.imwrite(f"/workspace/workspace/wumh/wuminghui/12_Smoke_fire/result/_{img_name}", img_bgr)

    
